import fitz  # PyMuPDF
import os
import shutil

# 输入和输出目录
input_dir = '../个体户'
output_dir = '../output_images'

# 创建输出目录如果它不存在
if not os.path.exists(output_dir):
    os.makedirs(output_dir)

# 支持的图像格式
image_extensions = {'.png', '.jpg', '.jpeg', '.bmp', '.tiff'}

for filename in os.listdir(input_dir):
    file_path = os.path.join(input_dir, filename)
    base_name, ext = os.path.splitext(filename)

    try:
        if ext.lower() == '.pdf':
            # 处理 PDF 文件
            doc = fitz.open(file_path)

            for page_num in range(len(doc)):
                page = doc.load_page(page_num)
                pix = page.get_pixmap()
                image_filename = f'{base_name}-{page_num + 1}.png'
                image_output_path = os.path.join(output_dir, image_filename)
                pix.save(image_output_path)

            print(f'Converted {filename} to images in {output_dir}')

        elif ext.lower() in image_extensions:
            # 处理图像文件
            image_output_path = os.path.join(output_dir, filename)
            shutil.copy2(file_path, image_output_path)
            print(f'Copied image {filename} to {image_output_path}')

    except Exception as e:
        print(f'Error processing {filename}: {e}')

print("All files have been processed.")
